from typing import Dict, Iterator
from app.jobs.pdfparser.internal.logic.splitter.llm.by_3steps_splitter import By3StepsSplitter
from app.jobs.pdfparser.internal.logic.splitter.llm.paper_abstract_extractor import PaperAbstractExtractor
from app.jobs.pdfparser.internal.logic.splitter.llm.paper_abstract_extractor_utils import merge_paper_basic_info, merge_other_research_fields
from app.jobs.pdfparser.internal.logic.splitter.text_splitter import TextSegment, TextSplitter
from app.jobs.pdfparser.internal.svc.servicecontext import ServiceContext
from langchain_openai import ChatOpenAI

class LlmTextSplitter(TextSplitter):
    def __init__(self, llm: ChatOpenAI, svcCtx: ServiceContext):
        super().__init__()
        self._abstract_extractor = PaperAbstractExtractor(llm)
        self._llm_splitter = By3StepsSplitter(llm)
        self._svcCtx = svcCtx
    
    def split(self, text: str, extra_kwargs: Dict = None) -> Iterator[TextSegment]:
        abstract_state = self._abstract_extractor.extract(text)
        my_summary = abstract_state.get("my_summary", {})
        
        user_id = extra_kwargs["user_id"]
        paper_file_id = extra_kwargs["paper_file_id"]
        paper_basic_info = merge_paper_basic_info(my_summary)
        abstract_summary = merge_other_research_fields(my_summary)
        self._svcCtx.paper_segment_model.insert_meta(user_id, paper_file_id, paper_basic_info, abstract_summary)
        
        for segment in self._llm_splitter.split_stream(abstract_state):
            yield TextSegment(
                        text=segment,
                        meta={},
                        splitter_id=self.splitter_id
                    )